"""
pypdf2 官方文档：https://pythonhosted.org/PyPDF2/
当前版本：1.26.0
目前更新：pypdf3，且更新日期为2年前
"""
import os
from PyPDF2 import PdfFileWriter, PdfFileReader


# 单文件加密
def pdf_encrypt(filename, password):
    """
    :param filename: 需要加密的文档
    :param password: 设置加密密码
    :return: 1，表示成功
    """
    pdf_reader = PdfFileReader(open(filename, 'rb'))
    pdf_writer = PdfFileWriter()
    for page in range(pdf_reader.getNumPages()):
        pdf_writer.addPage(pdf_reader.getPage(page))
    pdf_writer.encrypt(password)
    # 分隔文件名字，是个元祖类型
    new_file = os.path.splitext(filename)
    new_file = new_file[0] + '已加密' + '.pdf'
    # print('已经在此文档下生成加密的文档')
    with open(new_file, 'wb') as f:
        pdf_writer.write(f)
    return 1


# 解密
def pdf_decrypt(filename, password):
    """
    :param filename: 解密的文档
    :param password: 密码
    :return:
    """
    pdf_read = PdfFileReader(open(filename, 'rb'))
    if pdf_read.isEncrypted:
        try:
            pdf_read.decrypt(password)
            pdf_write = PdfFileWriter()
            for page in range(pdf_read.getNumPages()):
                pdf_write.addPage(pdf_read.getPage(page))
            new_file = os.path.splitext(filename)
            new_file = new_file[0] + '已解密' + '.pdf'
            # print('已经在此文档下生成解密的文档')
            with open(new_file, 'wb') as f:
                pdf_write.write(f)
            return 1
        except Exception as e:
            # 无法解密说明该模块无法解密高模板的加密软件
            # print("only algorithm code 1 and 2 are supported，无法解密该文档！")
            return 2
    else:
        # print("非加密程序")
        return 0


# 批量加密
def pdf_encrypts(filename_list, password):
    for filename in filename_list:
        pdf_encrypt(filename, password)
